import struct
import numpy as np
from plyfile import PlyData
from tqdm import tqdm

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def ply_to_splat(ply_file, splat_file):
    SH_C0 = 0.28209479177387814  # SH零阶常数

    # 读取 PLY
    plydata = PlyData.read(ply_file)
    vertices = plydata['vertex'].data

    # xyz (float32)
    xyz = np.vstack([vertices['x'], vertices['y'], vertices['z']]).T.astype(np.float32)

    # scale (exp 变换, float32)
    scales = np.vstack([
        vertices['scale_0'],
        vertices['scale_1'],
        vertices['scale_2']
    ]).T
    scales = np.exp(scales).astype(np.float32)

    # f_dc 转 rgb
    r = np.clip(0.5 + vertices['f_dc_0'] * SH_C0, 0, 1)
    g = np.clip(0.5 + vertices['f_dc_1'] * SH_C0, 0, 1)
    b = np.clip(0.5 + vertices['f_dc_2'] * SH_C0, 0, 1)
    rgb = (np.stack([r, g, b], axis=1) * 255).astype(np.uint8)

    # opacity (sigmoid + uint8)
    opacity = (sigmoid(vertices['opacity']) * 255).astype(np.uint8).reshape(-1, 1)

    # 合并成 RGBA
    colors = np.hstack([rgb, opacity])  # uint8[4]

    # rot → 先标准化，再映射到 uint8
    rots = np.vstack([
        vertices['rot_0'],
        vertices['rot_1'],
        vertices['rot_2'],
        vertices['rot_3']
    ]).T.astype(np.float32)

    # 标准化为单位四元数
    norms = np.linalg.norm(rots, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    rots_norm = rots / norms

    # 映射 [-1,1] → [0,255]
    rots_uint8 = (rots_norm * 128 + 128).clip(0, 255).astype(np.uint8)

    # 写入 SPLAT: xyz(float32), scale(float32x3), color(uint8x4), rot(uint8x4)
    with open(splat_file, "wb") as f:
        for i in tqdm(range(len(vertices))):
            # xyz
            f.write(struct.pack("fff", *xyz[i]))
            # scale
            f.write(struct.pack("fff", *scales[i]))
            # color (rgba)
            f.write(struct.pack("BBBB", *colors[i]))
            # rot
            f.write(struct.pack("BBBB", *rots_uint8[i]))

    print(f"✅ 转换完成: {len(vertices)} 个点写入 {splat_file}")


if __name__ == "__main__":
    ply_to_splat("data/ply/scene_mask_gs.ply", "data/splat/scene_mask_gs.splat")
